from ultralytics import YOLO
import pandas as pd

def train_model():
    model = YOLO("yolo11x-cls.pt")
    model.train(
        data='C:/Users/caruso/Desktop/YOLO_Project/images/',
        epochs=100,
        batch=16,
        imgsz=640,
        patience=10,
        project="agrilus",
        flipud=0.5,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.2,
        hsv_h=0.015, hsv_s=0.7, hsv_v=0.4,
        weight_decay=0.0005
    )


def predict_model():
    model = YOLO('C:/Users/caruso/Desktop/YOLO_Project/agrilus/train/weights/best.pt')
    source = "C:/Users/caruso/Desktop/YOLO_Project/images/test_random"

    model.predict(source=source, save=True, imgsz=640, conf=0.5)

def main():
    train_model()
    predict_model()

if __name__ == '__main__':
    main()